import numpy as np
import cv2
from typing import Tuple, Optional


class ImageTransformManager:
    """图像变换矩阵管理器"""

    def __init__(self):
        self.matrix_ab = None  # A到B的变换矩阵
        self.matrix_ac = None  # A到C的变换矩阵
        self.matrix_bc = None  # B到C的变换矩阵
        self.image_b_size = None  # 图片B的原始大小
        self.image_d_size = None  # 图片D的新大小

    def set_transforms(self, matrix_ab: np.ndarray, matrix_ac: np.ndarray,
                       image_b_size: Tuple[int, int]):
        """
        设置已知的变换矩阵

        Args:
            matrix_ab: A到B的变换矩阵 (3x3)
            matrix_ac: A到C的变换矩阵 (3x3)
            image_b_size: 图片B的大小 (width, height)
        """
        self.matrix_ab = matrix_ab.copy()
        self.matrix_ac = matrix_ac.copy()
        self.image_b_size = image_b_size

        # 计算B到C的变换矩阵: matrix_bc = matrix_ac * matrix_ab^(-1)
        matrix_ab_inv = np.linalg.inv(matrix_ab)
        self.matrix_bc = np.dot(matrix_ac, matrix_ab_inv)

        print("已设置变换矩阵:")
        print(f"A->B变换矩阵:\n{self.matrix_ab}")
        print(f"A->C变换矩阵:\n{self.matrix_ac}")
        print(f"B->C变换矩阵:\n{self.matrix_bc}")

    def calculate_size_transform(self, new_size: Tuple[int, int]) -> np.ndarray:
        """
        计算从原始B尺寸到新尺寸D的变换矩阵

        Args:
            new_size: 新的图片D大小 (width, height)

        Returns:
            变换矩阵 (3x3)
        """
        old_w, old_h = self.image_b_size
        new_w, new_h = new_size

        # 计算缩放比例
        scale_x = new_w / old_w
        scale_y = new_h / old_h

        # 创建缩放矩阵
        size_transform = np.array([
            [scale_x, 0, 0],
            [0, scale_y, 0],
            [0, 0, 1]
        ], dtype=np.float32)

        return size_transform

    def get_adjusted_transform(self, new_size: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取调整后的变换矩阵

        Args:
            new_size: 新的图片D大小 (width, height)

        Returns:
            (matrix_ad, matrix_bd): A到D的变换矩阵, B到D的变换矩阵
        """
        if self.matrix_ab is None or self.matrix_bc is None:
            raise ValueError("请先设置变换矩阵")

        self.image_d_size = new_size

        # 获取尺寸变换矩阵
        size_transform = self.calculate_size_transform(new_size)

        # 方法1: 直接基于B->C变换调整
        # 先将B变换到D的尺寸，然后应用原来的B->C操作
        matrix_bd_direct = np.dot(self.matrix_bc, size_transform)

        # 方法2: 通过A作为中介计算
        # A->D = (B->C) * (A->B) * size_transform^(-1)
        size_transform_inv = np.linalg.inv(size_transform)
        matrix_ad = np.dot(self.matrix_bc, np.dot(self.matrix_ab, size_transform_inv))

        print(f"\n调整图片B尺寸从 {self.image_b_size} 到 {new_size}")
        print(f"尺寸变换矩阵:\n{size_transform}")
        print(f"B->D变换矩阵:\n{matrix_bd_direct}")
        print(f"A->D变换矩阵:\n{matrix_ad}")

        return matrix_ad, matrix_bd_direct

    def apply_transform_to_image(self, image: np.ndarray, transform_matrix: np.ndarray,
                                 output_size: Tuple[int, int]) -> np.ndarray:
        """
        对图像应用变换矩阵

        Args:
            image: 输入图像
            transform_matrix: 变换矩阵 (3x3)
            output_size: 输出图像大小 (width, height)

        Returns:
            变换后的图像
        """
        # 使用仿射变换（取2x3矩阵）
        affine_matrix = transform_matrix[:2, :]

        transformed = cv2.warpAffine(
            image,
            affine_matrix,
            output_size,
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(0, 0, 0)
        )

        return transformed


def demo_usage():
    """演示用法"""
    # 创建变换管理器
    manager = ImageTransformManager()

    # 示例数据
    # A到B的变换矩阵（比如旋转+平移）
    matrix_ab = np.array([
        [0.8, -0.6, 100],
        [0.6, 0.8, 50],
        [0, 0, 1]
    ], dtype=np.float32)

    # A到C的变换矩阵（在B的基础上又做了缩放平移）
    matrix_ac = np.array([
        [1.2, -0.9, 150],
        [0.9, 1.2, 100],
        [0, 0, 1]
    ], dtype=np.float32)

    # 图片B的原始大小
    image_b_size = (800, 600)

    # 设置变换矩阵
    manager.set_transforms(matrix_ab, matrix_ac, image_b_size)

    # 新的图片D大小
    new_size = (1200, 900)  # 放大1.5倍

    # 获取调整后的变换矩阵
    matrix_ad, matrix_bd = manager.get_adjusted_transform(new_size)

    print("\n=== 使用示例 ===")
    print("假设你有图片A，想要:")
    print("1. 通过matrix_ad直接从A变换到D")
    print("2. 通过matrix_bd从调整尺寸后的B变换到D")

    return manager, matrix_ad, matrix_bd


def create_test_image(size: Tuple[int, int], text: str) -> np.ndarray:
    """创建测试图像"""
    image = np.zeros((size[1], size[0], 3), dtype=np.uint8)
    # 添加一些视觉标记
    cv2.rectangle(image, (10, 10), (size[0] - 10, size[1] - 10), (255, 255, 255), 2)
    cv2.putText(image, text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

    # 添加坐标系
    cv2.arrowedLine(image, (100, 100), (200, 100), (0, 0, 255), 3)  # X轴-红色
    cv2.arrowedLine(image, (100, 100), (100, 200), (0, 255, 0), 3)  # Y轴-绿色

    return image


def complete_example():
    """完整示例"""
    print("=== 完整变换示例 ===")

    # 创建测试图像A
    image_a = create_test_image((400, 300), "Image A")

    # 定义变换矩阵
    matrix_ab = np.array([
        [1.0, 0.2, 50],  # 轻微剪切 + 平移
        [-0.1, 1.0, 30],
        [0, 0, 1]
    ], dtype=np.float32)

    matrix_ac = np.array([
        [1.5, 0.3, 100],  # 更大的变换
        [-0.15, 1.2, 60],
        [0, 0, 1]
    ], dtype=np.float32)

    # 创建变换管理器
    manager = ImageTransformManager()
    image_b_size = (600, 450)  # B的大小
    manager.set_transforms(matrix_ab, matrix_ac, image_b_size)

    # 生成图像B（模拟）
    image_b = manager.apply_transform_to_image(image_a, matrix_ab, image_b_size)

    # 现在要调整B的大小到D
    new_size_d = (900, 675)  # 1.5倍放大
    matrix_ad, matrix_bd = manager.get_adjusted_transform(new_size_d)

    # 生成最终图像D
    image_d = manager.apply_transform_to_image(image_a, matrix_ad, new_size_d)

    print(f"\n成功生成变换:")
    print(f"图像A尺寸: {image_a.shape[:2][::-1]}")
    print(f"图像B尺寸: {image_b_size}")
    print(f"图像D尺寸: {new_size_d}")
    print(f"最终图像D的实际尺寸: {image_d.shape[:2][::-1]}")

    return {
        'manager': manager,
        'images': {'A': image_a, 'B': image_b, 'D': image_d},
        'matrices': {'AB': matrix_ab, 'AC': matrix_ac, 'AD': matrix_ad, 'BD': matrix_bd}
    }


if __name__ == "__main__":
    # 运行演示
    demo_usage()
    print("\n" + "=" * 50)
    complete_example()

